import numpy as np
from numpy import dot,inf
from math import sqrt,pow
from numba import jit
import matplotlib.pyplot as plt
import queue
import time
@jit
def mixf(a,b,mix):
    return b*mix+a*(1-mix)
@jit
def normalize(v):
    return v/np.linalg.norm(v)
@jit
def refract( v,   n,  m):
    refracted=np.zeros(3)
    if m[2]==inf:
        return False,refracted
    uv=normalize(v)
    dt = np.dot(uv, n);
    ni_over_nt=1/m[2]
    if dt>0:
        n=-n
        dt = -dt;
        ni_over_nt=m[2]
    discriminant = 1.0 - ni_over_nt * ni_over_nt * (1 - dt * dt);
    if (discriminant > 0):
        refracted = ni_over_nt * (uv - n * dt) - n * sqrt(discriminant);
        #print(refracted)
        return True,refracted;
    else:
        return False,refracted;
@jit
def reflect(v,n):
    v=normalize(v)
    return v-2*dot(n,v)*n
def circle(S,R):
    #x=np.linspace(0.,2*3.1415,100)
    y=R*np.sin(np.linspace(0.,2*3.1415,100))+S[1]
    x=R*np.cos(np.linspace(0.,2*3.1415,100))+S[0]
    return x,y
@jit
def intersect_sphere(O, D, S, R):
    # Return the distance from O to the intersection of the ray (O, D) with the
    # sphere (S, R), or +inf if there is no intersection.
    # O and S are 3D points, D (direction) is a normalized vector, R is a scalar.
    a = dot(D, D)
    OS =S-O
    b = dot(D, OS)
    c = dot(OS, OS) - R * R
    disc = b * b -  a * c
    if disc > 0 and b>0:
        rt = np.sqrt(disc)
        if b-rt>0 and b+rt>0:
            return b-rt
        if b-rt<0 and b+rt>0:
            return b+rt
    return np.inf
@jit
def get_normal(s,m):
    return normalize(m-s)
num=2
gx=320
gy=240
bm=np.zeros((gy,gx,3))
pos=np.zeros((num,4))
col=np.zeros((num,3))
mat=np.zeros((num,3))

pos[0,:]=[0,0,-20,4]
pos[1,:]=[0,-10004,-20,10000.0]
col[0,:]=[1.0,1,1]
col[1,:]=[0.5,0.2,0.2]
mat[0,:]=[1,0.,1.1]
mat[1,:]=[1,0.,inf]
#maxq=0
#@jit(nopython=True)
def tracer():
    maxq=0
    invwidth=1./gx
    invheight=1.0/gy
    asp=gx/gy
    fov=30
    angle=np.tan(np.pi*0.5*fov/180)
    rays = queue.Queue()
    for i in range(gy):
        for j in range(gx):
            xx=(2*(j+0.5)*invwidth-1)*angle*asp
            yy=(1-2*((i+0.5)*invheight))*angle #1
                
            D=normalize(np.array([xx,yy,-1]))
            O=np.array([0,0,0])
            tmp=np.zeros(3)
            tmpc=np.ones(3)
            color=np.array([0.6,0.6,0.9])#背景色
            #光线放入队列
            rays.put([O,D,tmpc,0])
            while (rays.qsize())>0:
                maxq=maxq if maxq> rays.qsize() else  rays.qsize()
                # 从队列取得光线
                O,D,temc,depth=rays.get()
                # 追踪这条光线
                t=inf
                # 求交点
                obj=0
                for k in range(num):
                    d=intersect_sphere(O, D,pos[k,0:3], pos[k,3])
                    if d<t:
                        t=d
                        obj=k
                if t==inf or depth >20:
                    tmp+=temc*color if depth <20 else 0
                    #break
                else:
                    M=O+D*t
                    N=get_normal(pos[obj,0:3],M)
                    np.array([0,0,0])
                    #color[:]=(1+N)/2  #col[obj,:]
                    tmpc*=col[obj,:]
                    # 反射
                    D0=reflect(D,N)
                    #折射
                    ref,DF=refract( D,   N, mat[obj])
                    if ref: #有折射
                        n=-N if dot(N,D)>0 else N
                        facingratio=-dot(D,n)
                        #print(facingratio,pow(1-facingratio,3))
                        fresneleffect=mixf(pow(1-facingratio,3),1,0.1) #b*mix+a(1-mix)
                        ior=1.1
                        eta=ior if dot(N,D)>0 else 1/ior
                        cosi=-dot(n,D)
                        k=1-eta*eta*(1-cosi*cosi)
                        DF=D*eta+n*(eta*cosi-sqrt(k))
                        O=M+dot(N,D)*N*0.001
                        rays.put([O,D0,tmpc*fresneleffect,depth+1]) #反射光线
                        O=M+dot(N,DF)*N*0.001
                        rays.put([O,DF,tmpc*(1-fresneleffect)*0.5,depth+1]) #折射光线
                        #D=DF
                    else:#全反射
                        D=D0
                        O=M+dot(N,D)*N*0.001
                        rays.put([O,D,tmpc,depth+1])
            bm[i,j,0]=min(tmp[0],1.0) 
            bm[i,j,1]=min(tmp[1],1.0) 
            bm[i,j,2]=min(tmp[2],1.0) 
    return maxq
            
t0=time.time()        
maxq=tracer()
print('maxq=',maxq,'time=',time.time()-t0)
plt.imshow(bm/np.max(bm))
plt.show()